# -*- coding: utf-8 -*-
"""
Created on  2019

@author: QW

"""

# coding: utf-8
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import time
'''
pso算法标准公式：
vi= w*vi+c1*r1*(xi-pbest)+c2*r2*(xi-gbest)
xi = xi+vi
'''
from sklearn.model_selection import train_test_split
from sklearn.linear_model import  LinearRegression
class PSO():
    def __init__(self, N, dim, step):
        self.w = 0.5                                # 惯性权重 一般在0-1之间
        self.c1 = self.c2 = 2                       # 学习因子,通常c1=c2=2
        self.r1 = 0.8
        self.r2 = 0.3
        self.N = N                                  # 粒子数量
        self.dim = dim                              # 搜索维度
        self.step = step                            # 迭代次数
        self.X = np.zeros((self.N, self.dim))       # 所有粒子的位置和速度
        self.V = np.zeros((self.N, self.dim))
        self.pbest = np.zeros((self.N, self.dim))   # 个体经历的最佳位置
        self.gbest = np.zeros((1, self.dim))        # 全局最佳位置
        self.p_fit = np.zeros(self.N)               # 每个个体的历史最佳适应值
        self.fit = 1000000000
        self.data = pd.read_csv(r"E:\机器学习\水质数据建模预测\data\去噪后的数据.csv", sep=',',encoding='gbk')
        y = self.data.iloc[:,1]
        x = self.data.iloc[:,2:]
        # 划分数据集
        train_x,test_x,train_y,test_y = train_test_split(x,y,test_size=0.3,random_state=12)
        self.train_x = train_x
        self.test_x = test_x
        self.train_y = train_y
        self.test_y  = test_y

    # 目标函数，粒子群选择的是最优波长点数 d1,d2,d3,d4,d5,d6代表所在列的数字
    def function(self,d1,d2,d3,d4,d5,d6):
        model = LinearRegression()
        train_x = self.train_x.iloc[:,[d1,d2,d3,d4,d5,d6]]
        model.fit(train_x, self.train_y)
        test_x = self.test_x.iloc[:,[d1,d2,d3,d4,d5,d6]]
        result = model.predict(test_x)
        loss = np.sqrt(0.5/len(result)* (np.square(self.test_y-result)).sum())
        return loss        # 求解均方根误差

    # 初始化粒子群
    def init_Population(self):
        for i in range(self.N):
            for j in range(self.dim):
                self.X[i][j] = np.random.randint(0, 1044)
                self.V[i][j] = np.random.randint(0, 10)
            self.pbest[i] = self.X[i]
            print(self.X[i][0], self.X[i][1],self.X[i][2],self.X[i][3],self.X[i][4],self.X[i][5])
            tmp = self.function(self.X[i][0], self.X[i][1],self.X[i][2],self.X[i][3],self.X[i][4],self.X[i][5])
            self.p_fit[i] = tmp
            if (tmp < self.fit):
                self.fit = tmp
                self.gbest = self.X[i]
    # 更新粒子位置
    def iterator(self):
        fitness = []
        for t in range(self.step):
            for i in range(self.N):                           # 更新gbest\pbest
                temp = self.function(self.X[i][0], self.X[i][1],self.X[i][2],self.X[i][3],self.X[i][4],self.X[i][5])
                if (temp < self.p_fit[i]):                    # 更新个体最优
                    self.p_fit[i] = temp
                    self.pbest[i] = self.X[i]
                    if (self.p_fit[i] < self.fit):            # 更新全局最优
                        self.gbest = self.X[i]
                        self.fit = self.p_fit[i]
            for i in range(self.N):
                V = self.w * self.V[i] + self.c1 * self.r1 * (self.pbest[i] - self.X[i]) + self.c2 * self.r2 * (
                            self.gbest - self.X[i])
                X = np.round(self.X[i] + V)
                print("粒子",X)
                if 0 < X[0]< 1043 and 0 < X[1]< 1043 and 0 < X[2]< 1043 and 0 < X[3]< 1043 and 0 < X[4]< 1043 and 0 < X[5]< 1043 :
                    self.V[i] = V
                    self.X[i] = X
            print("i:", self.X)
            fitness.append(self.fit)
            print("适应度:", self.fit)  # 输出最优值
            print("最优解:", self.gbest)
        return fitness

start_time = time.time()  # 开始计时
my_pso = PSO(N=100, dim=6, step=200)
my_pso.init_Population()
fitness = my_pso.iterator()
end_time = time.time()  # 结束计时
plt.figure(1)
plt.title("Figure1")
plt.xlabel("iterators", size=14)
plt.ylabel("fitness", size=14)
t = np.array([t for t in range(0, 200)])
fitness = np.array(fitness)
plt.plot(t, fitness, color='b')
plt.show()
data = pd.DataFrame(fitness)
data.to_csv("pso最优粒子目标函数值.csv")
print('用时：', end_time - start_time, '秒')

'''
适应度: 0.15117066708988808
最优解: [564. 612. 510. 110. 221. 436.]
用时： 146.26212573051453 秒
'''